# -*- coding: utf-8 -*-

from flask import Blueprint, jsonify, request, g, current_app
from flasgger import swag_from
from whoosh.index import create_in, open_dir
from whoosh.fields import Schema, TEXT, ID
from whoosh.qparser import QueryParser, FuzzyTermPlugin
from whoosh.analysis import Tokenizer, Token
from whoosh.index import EmptyIndexError, LockError
from docx import Document
import pdfplumber
import jieba
import pandas as pd
import chardet
import os
import time
from bson import ObjectId
from ..utils import make_response
from ..config import Config

bp = Blueprint('search_scratch', __name__)

class JiebaTokenizer(Tokenizer):
    def __call__(self, text, **kwargs):
        t = Token()
        for word in jieba.cut(text):
            t.text = word
            t.boost = 1.0
            t.pos = 0
            t.startchar = 0
            t.endchar = len(word)
            yield t

# 定义Whoosh索引的Schema
schema = Schema(content=TEXT(stored=True, analyzer=JiebaTokenizer()), path=ID(stored=True))

# 索引目录为 `scratchFile` 目录
scratch_indexdir = Config.indexdir_scratch_path
scratch_file_directory = Config.resource_local_scratch_path  # `scratchFile` 文件目录路径

def setup_scratch_file_index():
    """设置 `scratchFile` 目录的索引，如果不存在则创建，如果存在则更新"""
    try:
        if not os.path.exists(scratch_indexdir):
            os.mkdir(scratch_indexdir)
            ix = create_in(scratch_indexdir, schema)
            print("索引目录创建成功")
        else:
            ix = open_dir(scratch_indexdir)
            print("索引已存在，准备更新")

        existing_paths = {hit['path'] for hit in ix.searcher().documents()}
        retry_count = 0
        while retry_count < 3:
            try:
                with ix.writer() as writer:
                    add_files_to_index(scratch_file_directory, writer, existing_paths)
                break
            except LockError:
                retry_count += 1
                print(f"索引被锁定，等待重试... (尝试 {retry_count}/3)")
                time.sleep(1)  # 等待 1 秒后重试
        if retry_count == 3:
            print("无法获取索引锁定，操作失败。")
    except EmptyIndexError as e:
        print(f"索引文件不存在或损坏，正在重新创建索引: {str(e)}")
        ix = create_in(scratch_indexdir, schema)
        with ix.writer() as writer:
            add_files_to_index(scratch_file_directory, writer)
        print("索引重新建立成功")
    except Exception as e:
        print(f"未知错误: {str(e)}")
        raise e

    g.scratch_ix = ix  # 将索引对象放入全局对象g中

def add_files_to_index(directory, writer, existing_paths=None):
    """将指定目录下的文件添加到索引中，并显示进度"""
    if existing_paths is None:
        existing_paths = set()

    total_files = sum(len(files) for _, _, files in os.walk(directory))  # 总文件数
    processed_files = 0  # 已处理文件数

    for root, dirs, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            if file.startswith("~$") or file_path in existing_paths:
                continue  # 跳过临时文件或已存在的文件

            try:
                if file_path.endswith('.pdf'):
                    with pdfplumber.open(file_path) as pdf:
                        content = ''.join(page.extract_text() or '' for page in pdf.pages)
                    print(f"Adding PDF to index: {file_path}")
                    writer.add_document(content=content, path=file_path)
                elif file_path.endswith('.docx'):
                    content = extract_text_from_word(file_path)
                    print(f"Adding Word to index: {file_path}")
                    writer.add_document(content=content, path=file_path)
                elif file_path.endswith('.doc'):
                    content = extract_text_from_word(file_path)
                    print(f"Adding Word to index: {file_path}")
                    writer.add_document(content=content, path=file_path)
                elif file_path.endswith('.xlsx'):
                    content = extract_text_from_excel(file_path)
                    print(f"Adding Excel to index: {file_path}")
                    writer.add_document(content=content, path=file_path)
                elif file_path.endswith('.txt'):
                    content = read_file_with_encoding_detection(file_path)
                    print(f"Adding Text file to index: {file_path}")
                    writer.add_document(content=content, path=file_path)
                else:
                    print(f"Unsupported file format: {file_path}")
            except Exception as e:
                print(f"Failed to process file {file_path}: {e}")

            processed_files += 1
            progress = (processed_files / total_files) * 100
            print(f"索引进度: {progress:.2f}% ({processed_files}/{total_files})")  # 显示进度

def extract_text_from_word(docx_path):
    """从Word文件中提取文本"""
    try:
        doc = Document(docx_path)
        content = "\n".join([para.text for para in doc.paragraphs])
        return content
    except Exception as e:
        print(f"Error extracting Word file {docx_path}: {e}")
        return ""

def extract_text_from_excel(excel_path):
    """从Excel文件中提取文本"""
    try:
        content = ""
        xls = pd.ExcelFile(excel_path)
        for sheet_name in xls.sheet_names:
            df = pd.read_excel(excel_path, sheet_name=sheet_name)
            content += df.to_string(header=False, index=False)
        return content
    except Exception as e:
        print(f"Error extracting Excel file {excel_path}: {e}")
        return ""

def read_file_with_encoding_detection(file_path):
    """检测文件编码并读取文本内容"""
    try:
        with open(file_path, 'rb') as f:
            raw_data = f.read()
            result = chardet.detect(raw_data)
            encoding = result['encoding']
            if not encoding:
                encoding = 'utf-8'  # 默认使用 utf-8 编码
            return raw_data.decode(encoding, errors='ignore')
    except Exception as e:
        print(f"Error reading text file {file_path}: {e}")
        return ""

@bp.before_app_request
def before_request():
    """在每个请求之前设置 `scratchFile` 索引"""
    setup_scratch_file_index()

def preprocess_query(query):
    """对查询内容进行预处理"""
    words = jieba.cut(query)
    important_words = [word for word in words if word not in {"相关", "信息", "的", "是", "帮我查询", "搜索", "查询"}]
    return " ".join(important_words)

@bp.route('/search/scratch', methods=['POST'])
@swag_from({
    'parameters': [
        {
            'name': 'messages',
            'in': 'body',
            'required': True,
            'schema': {
                'type': 'object',
                'properties': {
                    'model': {
                        'type': 'string',
                        'example': 'llama3'
                    },
                    'messages': {
                        'type': 'array',
                        'items': {
                            'type': 'object',
                            'example': {"role": "file", "content": "搜索长沙元元熊餐饮管理有限公司"}
                        }
                    }
                }
            },
            'description': '输入消息数组，通常包含用户的查询内容'
        }
    ],
    'responses': {
        200: {
            'description': 'Successful Response',
            'examples': {
                'application/json': {
                    "results": [
                        {
                            "path": "D:\\Program Files\\AI学习文件1\\TEST\\example.pdf",
                            "snippet": "长沙元元熊餐饮管理有限公司...",
                            "metadata": {
                                "directory_name": "example_dir",
                                "uploader_id": "123456",
                                "timestamp": "2024-08-28 12:34:56"
                            }
                        }
                    ]
                }
            }
        }
    }
})
def search_in_scratch_files():
    """
    处理 `scratchFile` 文件的搜索请求
    ---
    tags:
      - 知识库相关接口
    """
    data = request.get_json()

    # 添加数据结构验证
    if not data or 'messages' not in data or not data['messages']:
        return make_response(message="Invalid input format: 'messages' is missing or empty", status=400)

    content = data['messages'][0].get('content')
    role = data['messages'][0].get('role')

    # 检查 role 和 content 是否存在
    if not content or not role:
        return make_response(message="Missing 'role' or 'content' in request data", status=400)

    try:
        processed_query = preprocess_query(content)
        print(f"Processed query: {processed_query}")  # 调试输出

        with g.scratch_ix.searcher() as searcher:
            parser = QueryParser("content", g.scratch_ix.schema)
            parser.add_plugin(FuzzyTermPlugin())
            query = parser.parse(processed_query)
            print(f"Parsed query: {query}")  # 调试输出

            results = searcher.search(query)
            response_data = []

            # MongoDB 数据库访问
            db = current_app.db  # 获取 MongoDB 数据库实例

            for result in results:
                path_parts = result['path'].split(Config.resource_scratch_path)
                if len(path_parts) > 1:
                    file_path = path_parts[1]  # 只有在split结果长度足够时才取第二个元素
                else:
                    file_path = result['path']  # 如果"resource"不在路径中，使用完整路径

                print(f"Found in file: {file_path}")  # 调试输出

                # 从 MongoDB 获取文件元数据
                file_metadata = db.files.find_one({"file_path": result['path']})

                if file_metadata:
                    response_data.append({
                        "path": file_path,
                        "snippet": result.highlights("content"),
                        "metadata": {
                            "directory_name": file_metadata.get("directory_name", ""),
                            "uploader_id": file_metadata.get("uploader_id", ""),
                            "timestamp": time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(file_metadata.get("timestamp", 0)))
                        }
                    })
                else:
                    response_data.append({
                        "path": file_path,
                        "snippet": result.highlights("content"),
                        "metadata": None
                    })

            if not results:
                print("No results found.")  # 调试输出

            return make_response(data=response_data[:3], message="Success", status=200)

    except Exception as e:
        print(f'Error occurred: {str(e)}')
        return make_response(message=str(e), status=500)

@bp.route('/search/file_by_remark', methods=['GET'])
@swag_from({
    'parameters': [
        {
            'name': 'remark',
            'in': 'query',
            'type': 'string',
            'required': True,
            'description': '文件备注名称'
        }
    ],
    'responses': {
        200: {
            'description': '文件查询成功',
            'schema': {
                'type': 'object',
                'properties': {
                    'files': {
                        'type': 'array',
                        'items': {
                            'type': 'object',
                            'properties': {
                                'file_path': {'type': 'string'},
                                'file_name': {'type': 'string'},
                                'timestamp': {'type': 'number'}
                            }
                        }
                    }
                }
            }
        }
    }
})
def search_file_by_remark():
    """
    根据备注名称查询文件地址 临时文件夹
    ---
    tags:
      - 知识库相关接口
    """
    remark = request.args.get('remark')

    if not remark:
        return make_response(message="缺少备注名称", status=400)

    db = current_app.db
    files = list(db.files.find({"remark": remark}))

    if not files:
        return make_response(message="未找到匹配的文件", status=404)

    response_data = [{"file_path": f["file_path"], "file_name": f["file_name"], "timestamp": f["timestamp"]} for f in files]

    return make_response(data=response_data, message="查询成功", status=200)
@bp.route('/update/remark_by_old', methods=['POST'])
@swag_from({
    'parameters': [
        {
            'name': 'old_remark',
            'in': 'formData',
            'type': 'string',
            'required': True,
            'description': '旧的备注名称'
        },
        {
            'name': 'new_remark',
            'in': 'formData',
            'type': 'string',
            'required': True,
            'description': '新的备注名称'
        }
    ],
    'responses': {
        200: {
            'description': '备注名称更新成功',
            'schema': {
                'type': 'object',
                'properties': {
                    'message': {'type': 'string'},
                    'updated_count': {'type': 'integer'}
                }
            }
        },
        400: {
            'description': '请求参数错误',
            'schema': {
                'type': 'object',
                'properties': {
                    'message': {'type': 'string'}
                }
            }
        },
        404: {
            'description': '未找到匹配的文件',
            'schema': {
                'type': 'object',
                'properties': {
                    'message': {'type': 'string'}
                }
            }
        }
    }
})
def update_remark_by_old():
    """
    根据旧的备注名称更新文件的备注名称
    ---
    tags:
      - 知识库相关接口
    """
    old_remark = request.form.get('old_remark')
    new_remark = request.form.get('new_remark')

    if not old_remark or not new_remark:
        return make_response(message="缺少旧的或新的备注名称", status=400)

    try:
        db = current_app.db
        result = db.files.update_many({"remark": old_remark}, {"$set": {"remark": new_remark}})

        if result.matched_count == 0:
            return make_response(message="未找到匹配的文件", status=404)

        return make_response(message="备注名称更新成功", data={"updated_count": result.modified_count}, status=200)

    except Exception as e:
        print(f"Error occurred while updating remark: {str(e)}")
        return make_response(message="服务器内部错误", status=500)

